package xcrypto

import (
	"crypto/ecdsa"
	"crypto/elliptic"
	"crypto/rand"
	"crypto/sha256"
	"math/big"

	"gitee.com/shuohe/xutils/xbuf"
	"gitee.com/shuohe/xutils/xerror"
)

var ECC_256_CURVE = elliptic.P256()
var ECC_256_BLOCK_SIZE = 32 // (ECC_256_CURVE.Params().BitSize + 7) >> 3
var eccZeroBlock [32]byte

//=============================================================================

/**
 * 包装过的 ECC 公钥对象。
 * @inherits IPublicKey, IEcPublicKey
 */
type EccPublicKey ecdsa.PublicKey

/**
 * @implements IKey.Algo
 */
func (this *EccPublicKey) Algo() string {
	return "ECC-256"
}

/**
 * @implements IKey.Type
 */
func (this *EccPublicKey) Type() string {
	return KeyType_Public
}

/**
 * @implements IKey.ExportX
 */
func (this *EccPublicKey) ExportToXBuf(xb *xbuf.Buffer) {
	xb.PutVStr(this.Algo())
	xb.PutVStr(this.Type())

	xbytes := this.X.Bytes()
	xb.PutBytes(eccZeroBlock[:32-len(xbytes)])
	xb.PutBytes(xbytes)

	ybytes := this.Y.Bytes()
	xb.PutBytes(eccZeroBlock[:32-len(ybytes)])
	xb.PutBytes(ybytes)
}

/**
 * @implements IKey.Export
 */
func (this *EccPublicKey) Export() []byte {
	xb := xbuf.New(2*32 + 32)
	this.ExportToXBuf(xb)

	return xb.Bytes()
}

/**
 * @implements IKey.ExportTo
 */
func (this *EccPublicKey) ExportTo(buf []byte) int {
	xb := xbuf.WrapForWrite(buf, false)
	this.ExportToXBuf(xb)

	return xb.Wi
}

/**
 * @implements IKey.Equals
 */
func (this *EccPublicKey) Equals(other IKey) bool {
	otherkey, ok := other.(*EccPublicKey)
	if !ok {
		return false
	}

	return this.X.Cmp(otherkey.X) == 0 && this.Y.Cmp(otherkey.Y) == 0
}

/**
 * @implements IPublicKey.Verify
 */
func (this *EccPublicKey) Verify(data []byte, signature []byte) bool {
	if len(signature) != 64 {
		return false
	}

	h := sha256.Sum256(data)
	r := new(big.Int).SetBytes(signature[:32])
	s := new(big.Int).SetBytes(signature[32:])

	return ecdsa.Verify((*ecdsa.PublicKey)(this), h[:], r, s)
}

//=============================================================================

/**
 * 包装过的私钥对象。
 * @inherits IPrivateKey, IEcPrivateKey
 */
type EccPrivateKey ecdsa.PrivateKey

/**
 * @implements IKey.Algo
 */
func (this *EccPrivateKey) Algo() string {
	return "ECC-256"
}

/**
 * @implements IKey.Type
 */
func (this *EccPrivateKey) Type() string {
	return KeyType_Private
}

/**
 * @implements IKey.ExportX
 */
func (this *EccPrivateKey) ExportToXBuf(xb *xbuf.Buffer) {
	xb.PutVStr(this.Algo())
	xb.PutVStr(this.Type())

	xbytes := this.X.Bytes()
	xb.PutBytes(eccZeroBlock[:ECC_256_BLOCK_SIZE-len(xbytes)])
	xb.PutBytes(xbytes)

	ybytes := this.Y.Bytes()
	xb.PutBytes(eccZeroBlock[:ECC_256_BLOCK_SIZE-len(ybytes)])
	xb.PutBytes(ybytes)

	dbytes := this.D.Bytes()
	xb.PutBytes(eccZeroBlock[:ECC_256_BLOCK_SIZE-len(dbytes)])
	xb.PutBytes(dbytes)
}

/**
 * @implements IKey.Export
 */
func (this *EccPrivateKey) Export() []byte {
	xb := xbuf.New(3*32 + 32)
	this.ExportToXBuf(xb)

	return xb.Bytes()
}

/**
 * @implements IKey.ExportTo
 */
func (this *EccPrivateKey) ExportTo(buf []byte) int {
	xb := xbuf.WrapForWrite(buf, false)
	this.ExportToXBuf(xb)

	return xb.Wi
}

/**
 * @implements IKey.Equals
 */
func (this *EccPrivateKey) Equals(other IKey) bool {
	otherkey, ok := other.(*EccPrivateKey)
	if !ok {
		return false
	}

	return this.X.Cmp(otherkey.X) == 0 && this.Y.Cmp(otherkey.Y) == 0 && this.D.Cmp(otherkey.D) == 0
}

/**
 * @implements IPrivateKey.Public
 */
func (this *EccPrivateKey) Public() IPublicKey {
	return (*EccPublicKey)(&this.PublicKey)
}

/**
 * @implements IPrivateKey.Sign
 */
func (this *EccPrivateKey) Sign(data []byte) []byte {
	r, s, _ := ecdsa.Sign(rand.Reader, (*ecdsa.PrivateKey)(this), Sha256(data))

	var result [64]byte

	rbytes := r.Bytes()
	copy(result[32-len(rbytes):], rbytes)

	sbytes := s.Bytes()
	copy(result[64-len(sbytes):], sbytes)

	return result[:]
}

/**
 * @implements IEcPrivateKey.RecoverKey
 */
func (this *EccPrivateKey) RecoverKey(pk IPublicKey, sk ISecretKey) ([]byte, error) {
	pkey, ok := pk.(*EccPublicKey)
	if !ok {
		return nil, xerror.New("Incorrect key algorithm.")
	}

	x, y := this.Curve.ScalarMult(pkey.X, pkey.Y, this.D.Bytes())
	xbytes := x.Bytes()
	ybytes := y.Bytes()

	var data [64]byte

	copy(data[32-len(xbytes):], xbytes)
	copy(data[2*32-len(ybytes):], ybytes)

	return sk.Encrypt(data[:]), nil
}

//=============================================================================

/**
 * 根据指定类型创建相应密钥对象。
 */
func EccKeyFactory(keyType string, keyData []byte) (IKey, int, error) {
	switch keyType {
	case KeyType_Public:
		// 生成一个公钥
		if len(keyData) < 2*32 {
			return nil, 0, errKeyData()
		}

		k := &EccPublicKey{
			Curve: ECC_256_CURVE,
			X:     new(big.Int).SetBytes(keyData[:32]),
			Y:     new(big.Int).SetBytes(keyData[32:64]),
		}

		return k, 64, nil

	case KeyType_Private:
		// 生成一个私钥
		if len(keyData) != 3*32 {
			return nil, 0, errKeyData()
		}
		k := &EccPrivateKey{
			PublicKey: ecdsa.PublicKey{
				Curve: ECC_256_CURVE,
				X:     new(big.Int).SetBytes(keyData[:32]),
				Y:     new(big.Int).SetBytes(keyData[32:64]),
			},
			D: new(big.Int).SetBytes(keyData[64:96]),
		}

		return k, 96, nil
	}

	return nil, 0, errKeyType(keyType)
}

/**
 * 用默认曲线设置生成 ECC 公私钥对。
 */
func EccKeyGenerator() (IKey, error) {
	sk, _ := ecdsa.GenerateKey(ECC_256_CURVE, rand.Reader)

	return (*EccPrivateKey)(sk), nil
}
